#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sun Jun  8 15:16:39 2025

@author: fenghongli
"""
#热力图
import os
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt

input_dir = 'garch_volatility'

vol_dict = {}

# 读取所有波动率文件
for file in os.listdir(input_dir):
    if file.endswith('_volatility.csv'):
        try:
            df = pd.read_csv(os.path.join(input_dir, file))
            ts_code = file.replace('_volatility.csv', '')
            df['trade_date'] = pd.to_datetime(df['trade_date'], format='%Y%m%d')
            df.set_index('trade_date', inplace=True)
            vol_dict[ts_code] = df['volatility']
        except Exception as e:
            print(f"读取 {file} 失败：{e}")

# 合并成一个 DataFrame（按日期对齐）
vol_df = pd.DataFrame(vol_dict).dropna(how='any')

# 计算相关性矩阵
corr_matrix = vol_df.corr()

# 绘图
plt.figure(figsize=(14, 12))
sns.heatmap(corr_matrix, cmap='coolwarm', annot=True, fmt=".2f", linewidths=0.5)
plt.title('Heat map of asset volatility correlation', fontsize=16)
plt.tight_layout()
plt.savefig('volatility_correlation_heatmap.png', dpi=300)
plt.show()


#多资产波动率对比图

#示例：建设银行 vs 万科A
bank_code = '601939.SH'     # 建设银行
re_code = '000002.SZ'       # 万科A

# 加载两个波动率数据
bank_df = pd.read_csv(f'garch_volatility/{bank_code}_volatility.csv')
re_df = pd.read_csv(f'garch_volatility/{re_code}_volatility.csv')

# 处理时间
bank_df['trade_date'] = pd.to_datetime(bank_df['trade_date'], format='%Y%m%d')
re_df['trade_date'] = pd.to_datetime(re_df['trade_date'], format='%Y%m%d')

# 合并数据
df_merged = pd.merge(bank_df, re_df, on='trade_date', how='inner', suffixes=('_bank', '_re'))

# 绘图
plt.figure(figsize=(12, 5))
plt.plot(df_merged['trade_date'], df_merged['volatility_bank'], label=f'{bank_code}（银行）', color='blue')
plt.plot(df_merged['trade_date'], df_merged['volatility_re'], label=f'{re_code}（房企）', color='green')
plt.title('Comparison of volatility between representative FI‘s and real estate companies')
plt.xlabel('Transaction Date')
plt.ylabel('Conditional volatility')
plt.legend()
plt.grid(True)
plt.tight_layout()
plt.savefig('bank_vs_re_volatility.png', dpi=300)
plt.show()
